package brainbots.brains;

import brainbots.bot.Body;
import brainbots.bot.BrainBot;
import ratson.utils.DoubleMatrix;
import ratson.utils.MatrixGenerator;

public class MatrixBrain implements Brain {

	private DoubleMatrix IX;
	private int sizeV;
	
	/**
	 * 
	 * @param sizeX - size of internal state vector
	 * @param sizeV - size of internal vecor, non-linearly transformation
	 * @param sizeU - size of input vector (must match body sensor count)
	 * @param sizeY - size of the output vector (must match body motor count)
	 */
	public MatrixBrain(int sizeU, int sizeY, int sizeX, int sizeV){
		
		initBrain(sizeU, sizeY, sizeX, sizeV);
	}
	
	/**Size of input data vector (number of sensors*/
	public int getSizeU() {
		return sizeU;
	}
	/**Size of the integrators couint*/
	public int getSizeX() {
		return sizeX;
	}
	/**size of the output vector (number of motors*/
	public int getSizeY() {
		return sizeY;
	}
	/**size of the internal non-linearly transformed vector*/
	public int getSizeV() {
		return sizeV;
	}
	
	/**Create brain, compatible with body (it does not stores body reference, uses only sizes)
	 * V - number of really distinguishable control signals
	 * X - number of internal state variables
	 * */
	public MatrixBrain(Body body, int sizeV, int sizeX){
		initBrain(body.getSensorCount(), body.getMotorCount(), sizeX, sizeV);
	}
	
	
	private void initBrain(int sizeU, int sizeY, int sizeX, int sizeV) {
		this.sizeX = sizeX;
		this.sizeY = sizeY;
		this.sizeU = sizeU;
		this.sizeV = sizeV;
		
		X_Y = new DoubleMatrix(1, sizeX+sizeY);
		V = new DoubleMatrix(1, sizeV);
		X_U = new DoubleMatrix(1, sizeX+sizeU);
		
		A1 = new DoubleMatrix(X_U.getWidth(), V.getWidth());
		A2 = new DoubleMatrix(V.getWidth(), X_Y.getWidth());
		
		IX = new DoubleMatrix(1, sizeX);
	}

	/**calculate static part of the brain*/
	private void calcDX(){
		//calculate the V vector
		V.mult(X_U, A1);
		
		//apply nonlinear transform
		for (int i = 0; i<V.getWidth(); ++i){
			V.set(0, i, nlfunc(V.get(0,i)));
		}
		//Apply second linear combination
		X_Y.mult(V, A2);
		
	}
	
	public DoubleMatrix A1, A2;
	private DoubleMatrix X_U, V, X_Y;
	private int sizeU;
	private int sizeX;
	private int sizeY;
	
	/**Nonlineafr transformation function. Should be fast and non-linear
	 * @param x
	 * @return
	 */
	private static double nlfunc(double x){
		return x/(Math.abs(x)+1)+0.1;
	}
	
	@Override
	public void simulate(double dt, Body bot) {
		
		if (bot.getSensorCount() != sizeU )
			throw new RuntimeException("U size and sensor count for the bot must be the same");
		//fill with body signals
		for(int i =0;i<sizeU;++i){
			X_U.set(0,i, bot.getSensor(i));
		}
		//fill with integrated value
		for(int i = 0;i<sizeX;++i){
			X_U.set(0,i+sizeU, IX.get(0, i));
		}
		
		//calculate the Y and DX
		calcDX();
		
		//updateintegrator value
		for(int i =0;i<sizeX;++i){
			IX.set(0,i,IX.get(0, i) + dt*X_Y.get(0, i));
		}
		
		//fill the motor controls
		for(int i =0; i<sizeY;++i){
			bot.setMotor(i, limitFunc(X_Y.get(0, sizeX+i)));
		}
	}
	
	private double limitFunc(double x){
		if (x>0) return 1;
		else return 0;
		/*if (x>1) return 1;
		if (x<-1) return -1;
		return x;
		*/
	}
	
	
	public void randomInit(){
		MatrixGenerator g = new MatrixGenerator(){
			@Override
			public double get(int y, int x) {
				return Math.random()*2-1;
			}
		};
		A1.fill(g);
		A2.fill(g);
		IX.fill(0.0);
	}

	
}
